import cv2
import pytesseract
import os
# 指定 Tesseract 路径
pytesseract.pytesseract.tesseract_cmd = r'D:\Tesseract-OCR\tesseract.exe'
# 指定 TESSDATA_PREFIX 环境变量，使用原始字符串
os.environ['TESSDATA_PREFIX'] = r'D:\Tesseract-OCR\tessdata'

def remove_text_from_image(image_path):
  # 读取图片
  image = cv2.imread(image_path)
  # 将图片转换为灰度图
  gray = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY)
  # 使用 pytesseract 进行文字识别，设置语言为简体中文
  data = pytesseract.image_to_data(gray, lang = 'chi_sim', output_type = pytesseract.Output.DICT)
  for i in range(len(data['text'])):
    # 去除前后空格并检查是否有文字
    if data['text'][i].strip() and int(data['conf'][i]) > 20:
      (x, y, w, h) = (data['left'][i], data['top'][i], data['width'][i], data['height'][i])
      # 对文字区域进行裁剪
      roi = image[y:y + h, x:x + w]
      # 检查 roi 是否为空
      if roi.size > 0:
        # 对文字区域进行模糊处理
        blurred = cv2.GaussianBlur(roi, (25, 25), 0)
        image[y:y + h, x:x + w] = blurred
  return image

def process_images_in_directory(directory):
  for filename in os.listdir(directory):
    if filename.lower().endswith(('.png', '.jpg', '.jpeg')):
      image_path = os.path.join(directory, filename)
      processed_image = remove_text_from_image(image_path)
      # 生成新的文件名
      new_filename = os.path.splitext(filename)[0] + '_1' + os.path.splitext(filename)[1]
      new_image_path = os.path.join(directory, new_filename)
      # 保存处理后的图片
      cv2.imwrite(new_image_path, processed_image)
      print(f"Processed and saved: {new_image_path}")

if __name__ == "__main__":
  # 请将此处替换为你的图片目录路径
  directory = r'D:\NodeTest\images\demo'
  process_images_in_directory(directory)
